import os
import requests
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('spark-api-client')

# 星火大模型配置（可通过环境变量或直接修改）
SPARK_API_KEY = os.environ.get('SPARK_API_KEY',"Bearer WAFmivwUnPFfnSDdUHyi:BsSnFeAafCFMVzIbJjqR")
SPARK_API_URL = os.environ.get('SPARK_API_URL',"https://spark-api-open.xf-yun.com/v1/chat/completions")
DEFAULT_MODEL = "4.0Ultra"  # 默认模型版本


class SparkAPIError(Exception):
    """星火API调用异常基类"""
    pass


class SparkRequestError(SparkAPIError):
    """请求发送异常"""
    pass


class SparkResponseError(SparkAPIError):
    """响应处理异常"""
    pass


def call_spark_model(
        messages,
        model=DEFAULT_MODEL,
        stream=False,
        web_search_enable=True,
        web_search_mode="deep"
):
    """
    调用讯飞星火大模型API

    参数:
        messages: 对话历史列表，格式为[{"role": "user/assistant", "content": "对话内容"}, ...]
        model: 模型版本，默认"4.0Ultra"
        stream: 是否流式返回，默认False
        web_search_enable: 是否启用web搜索，默认True
        web_search_mode: 搜索模式，"deep"或"simple"，默认"deep"

    返回:
        模型响应内容（str），流式返回时为生成器

    异常:
        SparkRequestError: 请求发送失败
        SparkResponseError: 响应解析失败
    """
    # 构造请求头
    headers = {
        "Authorization": SPARK_API_KEY,
        "Content-Type": "application/json"
    }

    # 构造请求体
    payload = {
        "model": model,
        "user": "default_user",  # 可根据实际用户ID替换
        "messages": messages,
        "stream": stream,
        "tools": [
            {
                "type": "web_search",
                "web_search": {
                    "enable": web_search_enable,
                    "search_mode": web_search_mode
                }
            }
        ]
    }

    try:
        # 发送请求
        response = requests.post(
            url=SPARK_API_URL,
            json=payload,
            headers=headers,
            timeout=30,  # 30秒超时
            stream=stream  # 流式请求需要开启stream参数
        )
        response.raise_for_status()  # 触发HTTP错误（4xx/5xx）

    except requests.exceptions.RequestException as e:
        logger.error(f"请求星火API失败: {str(e)}")
        raise SparkRequestError(f"请求发送失败: {str(e)}") from e

    # 处理非流式响应
    if not stream:
        try:
            result = response.json()
            # 提取响应内容
            if "choices" in result and len(result["choices"]) > 0:
                return result["choices"][0]["message"]["content"]
            else:
                raise SparkResponseError(f"响应格式异常: 缺少choices字段 - {result}")
        except (KeyError, ValueError) as e:
            logger.error(f"解析星火API响应失败: {str(e)}，响应内容: {response.text}")
            raise SparkResponseError(f"响应解析失败: {str(e)}") from e

    # 处理流式响应（生成器）
    else:
        def stream_generator():
            for line in response.iter_lines():
                if line:
                    # 处理SSE格式（移除"data: "前缀）
                    line_str = line.decode('utf-8').lstrip('data: ')
                    if line_str == '[DONE]':
                        break
                    try:
                        chunk = json.loads(line_str)
                        if "choices" in chunk and len(chunk["choices"]) > 0:
                            content = chunk["choices"][0]["message"].get("content", "")
                            if content:
                                yield content
                    except (KeyError, ValueError) as e:
                        logger.warning(f"解析流式响应块失败: {str(e)}，块内容: {line_str}")
                        continue

        return stream_generator()


# 示例使用
if __name__ == "__main__":
    # 示例对话历史
    sample_messages = [
        {"role": "user", "content": "什么是人工智能？"},
        {"role": "assistant",
         "content": "人工智能是研究、开发用于模拟、延伸和扩展人的智能的理论、方法、技术及应用系统的一门新的技术科学。"},
        {"role": "user", "content": "它有哪些实际应用场景？"}
    ]

    try:
        # 调用大模型
        print("正在调用星火大模型...")
        response = call_spark_model(sample_messages)
        print("\n模型响应:")
        print(response)

    except SparkAPIError as e:
        print(f"\n调用失败: {e}")